use std::collections::HashMap;

use indexmap::IndexMap;

use bc_base::{
    TransactionId,
    transaction::{Transaction, TransactionInPoint, TransactionOutPoint},
};

use crate::transaction::TransactionConsensusError;

#[derive(Debug)]
pub(crate) struct TransactionMemPool {
    tx_pool: IndexMap<TransactionId, (Transaction, Vec<u8>)>, // mapTransactions  // (tx, tx_serialize)
    // UTXO(Unspent Transaction Output) used by the transactions in the pool
    used_utxos_by_txs_pool: HashMap<TransactionOutPoint, TransactionInPoint>, // mapNextTx

    tx_serialize_size_in_tx_pool: usize,
}

impl TransactionMemPool {
    pub fn new() -> Self {
        Self {
            tx_pool: IndexMap::new(),
            used_utxos_by_txs_pool: HashMap::new(),

            tx_serialize_size_in_tx_pool: 0,
        }
    }

    pub fn contains(&self, tx_id: &TransactionId) -> bool {
        self.tx_pool.contains_key(tx_id)
    }

    pub fn get_tx_from_pool(&self, tx_id: &TransactionId) -> Option<&(Transaction, Vec<u8>)> {
        self.tx_pool.get(tx_id)
    }

    pub fn get_tx_in_point(&self, out_point: &TransactionOutPoint) -> Option<&TransactionInPoint> {
        self.used_utxos_by_txs_pool.get(out_point)
    }
}

impl TransactionMemPool {
    pub fn add_tx_to_memory_pool(
        &mut self,
        tx_id: TransactionId,
        tx: Transaction,
        tx_serialize: Vec<u8>,
    ) -> Result<(), TransactionConsensusError> {
        for (index, tx_in) in tx.get_inputs().iter().enumerate() {
            let tx_in_point = TransactionInPoint::new(tx_id, index as u32);
            self.used_utxos_by_txs_pool
                .insert(*tx_in.get_prev_out_point(), tx_in_point);
        }

        self.tx_serialize_size_in_tx_pool += tx_serialize.len();
        let r = self.tx_pool.insert(tx_id, (tx, tx_serialize));
        assert_eq!(r, None);

        Ok(())
    }

    pub fn remove_tx_from_tx_pool(
        &mut self,
        tx_id: &TransactionId,
    ) -> Result<(), TransactionConsensusError> {
        match self.tx_pool.shift_remove(tx_id) {
            None => return Err(TransactionConsensusError::TransactionNotFoundInPool),
            Some((tx, tx_serialize)) => {
                for tx_in in tx.get_inputs() {
                    self.used_utxos_by_txs_pool
                        .remove(tx_in.get_prev_out_point());
                }
                self.tx_serialize_size_in_tx_pool -= tx_serialize.len();
            }
        }

        Ok(())
    }
}
